"""
bm25_search.py 是一个基于 BM25 算法和 Softmax 归一化的文本检索模块，用于从问题库中检索与查询最匹配的答案。
模块结合 Redis 缓存和 MySQL 数据库，支持问题加载、分词、BM25 评分、Softmax 归一化，并记录操作日志。
"""
import numpy as np
from rank_bm25 import BM25Okapi

from mysql_qa.db.mysql_client import MysqlClient
from mysql_qa.utils.preprocess import preprocess

from base.config import get_config
from base.logger import get_logger
from mysql_qa.cache.redis_client import RedisClient


# 导包


"""
代码实现思路：
1：加载：数据准备部分，基于mysql，redis中的数据，以及bm25算法，完成数据的加载和算法的初始化
    1.1：判断redis中是否有数据
    1.2：如果没有数据，则从mysql中加载数据，并写入到redis中
    1.3：如果有数据，则从redis中加载数据，
    1.4：完成bm25算法的初始化
2：检索：接收用户查询，基于bm25算法，返回最匹配的答案
    2.1：根据用户的query查询redis数据库
    2.2：如果有匹配的数据直接返回答案
    2.3：如果没有匹配的数据，则对用户查询进行分词
    2.4：基于bm25算法，计算查询与每个问题的相似度得分
    2.5：基于softmax归一化，将得分转换为概率分布
    2.6：判断概率最高的问题是否大于0.85
        2.6.1：如果大于0.85
            2.6.1.1：根据答案再次查找redis问题对应的答案
            2.6.1.2：如果有匹配的数据，则返回答案
            2.6.1.3：如果没有匹配的数据，则再从mysql中查询一次，返回答案
            2.6.1.4：将查询到的结果写入到redis中
        2.6.2：如果小于等于0.85
            2.6.2.1：执行RAG_QA的流程
"""
class BM25Search:
    def __init__(self):
        # 初始化日志和配置信息
        self.logger = get_logger()
        self.config = get_config()

        # 初始化数据库和缓存客户端
        self.mysql_client = MysqlClient()
        self.redis_client = RedisClient()

        # 初始化bm25算法
        self.bm25 = None

        # 初始化原始问题列表
        self.original_questions = None
        self.questions = None

        self.load_data()

    def load_data(self):
        """
        1：加载：数据准备部分，基于mysql，redis中的数据，以及bm25算法，完成数据的加载和算法的初始化
            1.1：判断redis中是否有数据
            1.2：如果没有数据，则从mysql中加载数据，并写入到redis中
            1.3：如果有数据，则从redis中加载数据，
            1.4：完成bm25算法的初始化
        :return: None
        """
        # 定义所有问题和原始问题的redis的key
        original_key = "qa_original_questions"
        tokenized_key = "qa_tokenized_questions"

        # 1.1：判断redis中是否有数据
        self.original_questions =  self.redis_client.get_data(original_key)
        self.tokenized_questions =  self.redis_client.get_data(tokenized_key)

        # 1.2：如果没有数据，则从mysql中加载数据，并写入到redis中
        if not self.original_questions or not self.tokenized_questions:
            self.logger.info("redis中数据不完整，执行初始化操作，从mysql中加载数据...")
            all_question = self.mysql_client.fetch_all_questions()
            print(all_question)
            questions = [question[0] for question in all_question]
            print(questions)

            self.tokenized_questions = [preprocess(question[0]) for question in all_question]

            # 将所有原始问题写入到redis中
            self.redis_client.set_data(original_key, questions)
            self.redis_client.set_data(tokenized_key, self.tokenized_questions)

        # 将分词后的问题列表赋值给属性
        self.questions = self.tokenized_questions

        # 1.4：完成bm25算法的初始化
        self.bm25 = BM25Okapi(self.questions)
        self.logger.info("初始化完成...")

    def softmax(self, scores):
        """
        基于softmax归一化，将得分转换为概率分布
        :param scores:
        :return:
        """
        exp_scores = np.exp(scores - np.max(scores))
        sum_exp_scores = np.sum(exp_scores)
        prob_dist = exp_scores / sum_exp_scores
        return prob_dist

    def search(self, query: str, threshold=0.85):
            """
            2：检索：接收用户查询，基于bm25算法，返回最匹配的答案
                2.1：根据用户的query查询redis数据库
                2.2：如果有匹配的数据直接返回答案
                2.3：如果没有匹配的数据，则对用户查询进行分词
                2.4：基于bm25算法，计算查询与每个问题的相似度得分
                2.5：基于softmax归一化，将得分转换为概率分布
                2.6：判断概率最高问题是否大于0.85
                    2.6.1：如果大于0.85
                        2.6.1.1：根据答案再次查找redis问题对应的答案
                        2.6.1.2：如果有匹配的数据，则返回答案
                        2.6.1.3：如果没有匹配的数据，则再从mysql中查询一次，返回答案
                        2.6.1.4：将查询到的结果写入到redis中
                    2.6.2：如果小于等于0.85
                        2.6.2.1：执行RAG_QA的流程
            """
            # 检查用户输入是否为空
            if not query.strip():
                self.logger.error("用户非法输入，请重新输入...")
                return None, False

            # 2.1：根据用户的query查询redis数据库
            cached_answer = self.redis_client.get_answer(query)

            # 2.2：如果有匹配的数据直接返回答案
            if cached_answer:
                self.logger.info("查询结果已缓存，直接返回结果...")
                return cached_answer, False

            # 2.3：如果没有匹配的数据，则对用户查询进行分词
            query_tokens = preprocess(query)

            # 2.4：基于bm25算法，计算查询与每个问题的相似度得分
            scores = self.bm25.get_scores(query_tokens)

            # 2.5：基于softmax归一化，将得分转换为概率分布
            softmax_scores = self.softmax(scores)

            # 找到得分最高的索引和对应的得分
            max_index = np.argmax(softmax_scores)
            max_score = softmax_scores[max_index]

            # 2.6：判断概率最高问题是否大于0.85
            if max_score > threshold:
                # 得到与原始问题最匹配的其他问题
                original_question = self.original_questions[max_index]

                # 2.6.1.1：根据答案再次查找redis问题对应的答案
                cached_answer = self.redis_client.get_answer(original_question)
                if cached_answer:
                    self.logger.info("查询结果已缓存，直接返回结果...")
                    return cached_answer, False
                else:
                    self.logger.info("查询结果未缓存，直接读取mysql返回结果...")
                    # 2.6.1.3：如果没有匹配的数据，则再从mysql中查询一次，返回答案
                    cached_answer = self.mysql_client.fetch_answer(original_question)
                    # 对提问的问题可以设置过期时间，具体设置多久跟业务相关，最简单的方式可以粗暴设置为固定时间，如1天
                    self.redis_client.set_answer(original_question, cached_answer, expire_time=60 * 60 * 24)
                    return cached_answer, False

            # 2.6.2：如果小于等于0.85
            return None, True

if __name__ == '__main__':
    bm25_search = BM25Search()
    # bm25_search.init()
    bm25_search.search("windows如何安装redis")
